import matplotlib.pyplot as plt
import matplotlib.patches as patches
import numpy as np


def plot_matern_cluster_process(
    ax,
    points,
    center,
    radius_w,
    in_circle,
    cluster_sets_for_draw,
    cluster_sets,
    parent_points_x,
    parent_points_y,
    num_clusters,
    radius,  # 添加 radius 作为参数
):
    circle1 = plt.Circle(center, radius_w, fill=False, color="black")
    ax.add_patch(circle1)
    ax.scatter(points[:, 0], points[:, 1], label="All Points")
    parent_points = np.array(
        [(parent_points_x[i], parent_points_y[i]) for i in range(num_clusters)]
    )
    ax.scatter(
        parent_points[:, 0],
        parent_points[:, 1],
        marker="^",
        s=50,
        c="green",
        label="Cluster Centers",
    )
    ax.set_xlim(center[0] - radius, center[0] + radius)
    ax.set_ylim(center[1] - radius, center[1] + radius)
    ax.set_aspect("equal", adjustable="box")
    ax.set_title("Matern Cluster Process")
    ax.set_xlabel("X-axis")
    ax.set_ylabel("Y-axis")
    ax.legend()
    ax.scatter(points[in_circle, 0], points[in_circle, 1], c="red")
    for cluster_label, cluster_points in cluster_sets_for_draw.items():
        if cluster_label != -1:
            ax.scatter(
                [p[0] for p in cluster_points],
                [p[1] for p in cluster_points],
                label=f"Cluster {cluster_label}",
            )
        else:
            ax.scatter(
                [p[0] for p in cluster_points],
                [p[1] for p in cluster_points],
                label="Noise Points",
                marker="x",
                c="gray",
            )
    for cluster_label, cluster_points in cluster_sets.items():
        if cluster_label != -1:
            square_indices = cluster_points["square_indices"]
            if len(square_indices) > 0:
                square_points = points[square_indices]
                center_x = np.mean(square_points[:, 0])
                center_y = np.mean(square_points[:, 1])
                center_tuple = (center_x, center_y)
                distances = np.sqrt(
                    (square_points[:, 0] - center_x) ** 2
                    + (square_points[:, 1] - center_y) ** 2
                )
                radius = np.max(distances)
                circle1 = plt.Circle(center_tuple, radius, fill=False, color="blue")
                ax.add_patch(circle1)
                for x, y in square_points:
                    rect = patches.Rectangle(
                        (x - 0.5, y - 0.5),
                        1,
                        1,
                        linewidth=1,
                        edgecolor="blue",
                        facecolor="none",
                    )
                    ax.add_patch(rect)
            no_square_indices = cluster_points["no_square_indices"]
            if len(no_square_indices) > 0:
                no_square_points = points[no_square_indices]
                center_x = np.mean(no_square_points[:, 0])
                center_y = np.mean(no_square_points[:, 1])
                center_tuple = (center_x, center_y)
                distances = np.sqrt(
                    (no_square_points[:, 0] - center_x) ** 2
                    + (no_square_points[:, 1] - center_y) ** 2
                )
                radius = np.max(distances)
                circle1 = plt.Circle(center_tuple, radius, fill=False, color="orange")
                ax.add_patch(circle1)
    ax.scatter(0, 0, marker="+", s=100, c="r", label="Willie")
    return ax
